'''
功能：负责参数的配置。
参数：-p 模板 -i 输入源 -o 输出目录 -e(参考的示例) --tasks [任务列表] --llm AzureLLM或其他 （先硬编码数据、llm、解析）
输出：jsonl或txt
使用yaml配置参数，然后用hydra加载参数并可以命令行覆盖
1 IO参数： 输入文件、输出文件、模板 √
2 任务相关  √
3 llm参数：配置、生成参数 model_name=gpt35, max_tokens,MAX_API_RETRY = 5  REQ_TIME_GAP = 10  √

【先硬编码task，再动态替换llm或解析器。其实任务和输入和输出解析相关的，只有llm可以灵活替换】
'''
import argparse


def add_llm_args(parser):
    # 1 配置
    group = parser.add_argument_group("LLM")
    group.add_argument(
        "--engine",
        type=str,
        default="GPT_TURBO",
        help="deployment name",
    )

    group.add_argument(
        "--api-type",
        type=str,
        default="Azure",
        help="deployment name",
    )

    group.add_argument(
        "--api-version",
        type=str,
        default="2023-05-15",
        help="api-version",
    )

    group.add_argument(
        "--api-base",
        type=str,
        default="",
        help="api base",
    )

    group.add_argument(
        "--api-key",
        type=str,
        default="",
        help="api key",
    )

    # 2 生成
    # fmt: off
    parser.add_argument("--tozh",
                        action="store_true",
                        help="Weather translate result to chinese.")
    group.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="temperature",
    )

    group.add_argument(
        "--topp",
        type=float,
        default=0.95,
        help="temperature",
    )

    group.add_argument(
        "--max-tokens",
        type=int,
        default=1024,
        help="maximum number of tokens produced in the output",
    )

    group.add_argument(
        "--frequency_penalty",
        type=float,
        default=0,
        help="Penalizes repeated tokens according to frequency.",
    )

    group.add_argument(
        "--presence_penalty",
        type=float,
        default=0,
        help="maximum number of tokens produced in the output",
    )

    # group.add_argument(
    #     "--max-api-retry",
    #     type=int,
    #     default=5,
    #     help="maximum api retry",
    # )
    #
    # group.add_argument(
    #     "--req-time-gap",
    #     type=int,
    #     default=10,
    #     help="REQ_TIME_GAP",
    # )
    # fmt: on
    return parser


def add_task_args(parser):
    group = parser.add_argument_group("TASK")
    group.add_argument(
        "--batch-size",
        type=int,
        default=1,
        help="batch_size",
    )

    group.add_argument(
        "--num-procs",
        type=int,
        default=1,
        help="batch_size",
    )

    group.add_argument(
        "--sleep-sec",
        type=int,
        default=5,
        help="batch_size",
    )

    parser.add_argument("--outfile", default="result.txt",
                        help="output file")

    parser.add_argument("--sep", default="\n\n",
                        help="sample seperator")
    return parser


def get_parser():
    parser = argparse.ArgumentParser(allow_abbrev=False)
    # io
    parser.add_argument("-i", "--input", default=None, required=True,
                        help="input data source")
    parser.add_argument("-o", "--outdir", default=None, required=True,
                        help="output directory")
    parser.add_argument("-p", "--prompt", default=None, required=True,
                        help="prompt template")
    parser.add_argument("-c", "--conf", default=None,
                        help="yaml file path.")  # 修改默认配置

    # parser.add_argument("-e", "--example", default=None,required=True,
    #                    help="prompt examples")

    # task列表
    parser.add_argument(
        "--tasks", required=True,
        nargs="+", default=[],
        # choices=INFER_REGISTRY.keys(),
    )

    # llm
    # gpt-35-turbo-16k' or 'GPT_TURBO' gpt-3.5-turbo
    # parser.add_argument("--llm", default="GPT_TURBO",
    #                    help="llm model name")
    # add args
    parser = add_llm_args(parser)
    parser = add_task_args(parser)

    return parser
